package com.linkedin.metadata.models.extractor;

import com.datahub.util.ModelUtils;
import com.linkedin.data.schema.RecordDataSchema;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.metadata.models.FieldSpec;
import com.linkedin.metadata.models.annotation.AspectAnnotation;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;

/** Extracts fields from a RecordTemplate based on the appropriate {@link FieldSpec}. */
@Slf4j
public class AspectExtractor {

  private AspectExtractor() {}

  public static Map<String, RecordTemplate> extractAspectRecords(RecordTemplate snapshot) {
    return ModelUtils.getAspectsFromSnapshot(snapshot).stream()
        .collect(
            Collectors.toMap(
                record -> getAspectNameFromSchema(record.schema()), Function.identity()));
  }

  private static String getAspectNameFromSchema(final RecordDataSchema aspectSchema) {
    final Object aspectAnnotationObj =
        aspectSchema.getProperties().get(AspectAnnotation.ANNOTATION_NAME);
    if (aspectAnnotationObj != null) {
      return AspectAnnotation.fromSchemaProperty(aspectAnnotationObj, aspectSchema.getFullName())
          .getName();
    }
    log.error(
        String.format(
            "Failed to extract aspect name from provided schema %s", aspectSchema.getName()));
    throw new IllegalArgumentException(
        String.format(
            "Failed to extract aspect name from provided schema %s", aspectSchema.getName()));
  }
}
